from . import bp

from flask import request,jsonify,make_response,current_app
from app.modules import Payload,Category,Session,Project,Record
from functools import wraps
from app.exceptions import MsgError
from app.parsers import urlencode,urldecode
import base64
from .websocket import update_base_info


# 捕捉异常
def catch_exception(func):
    @wraps(func)
    def wrapper(*args,**kw):
        try:
            ret = func(*args,**kw)
        except:
            pass
    return wrapper

# 登录状态验证,只加在路由上
def verify(func):
    @wraps(func)
    def wrapper(*args,**kw):
        cookie = request.cookies.get(current_app.config.get('SESSION_NAME'))
        if cookie is None:
            return jsonify({'code': 1,'data':{'msg':'登录失效！'}})
        if not Session.login_check(cookie):
            return jsonify({'code': 1,'data':{'msg':'登录失效！'}})
        return func(*args,**kw)
    return wrapper

# 1. 登录请求

@bp.route('/login/<password>',methods=['GET'])
def login(password):
    ret = {'code': 0, 'data': {'msg': '登录成功！'}}
    try:
        password = urldecode(password)
        session = Session.create_session(password)
        ret = make_response(jsonify(ret))
        ret.set_cookie(current_app.config.get('SESSION_NAME'),session.login_id,max_age=60*60*24,httponly=True)
        return ret
    except MsgError as e:
        ret['code'] = 1
        ret['data']['msg'] = e.msg
    
    return jsonify(ret)


# 2. 获取所有payload基本信息

@bp.route('/payload/list/all',methods=['GET'])
@verify
def get_all_payload():
    ret = {'code': 0,'data':{'count': 0,'payload_list':[]}}
    payload_list = Payload.get_all_payload_with_json()
    ret['data']['count'] = len(payload_list)
    ret['data']['payload_list'] = payload_list
    return jsonify(ret)

# 3. 获取指定payload详细信息

@bp.route('/payload/<payload_id>/get',methods=['GET'])
@verify
def get_payload_info(payload_id):
    ret = {'code': 0, 'data': {}}
    try:
        payload_id = int(payload_id)
        info = Payload.get_payload_info_with_json(payload_id)
        ret['data'] = info
        return jsonify(ret)
    except ValueError:
        ret['code'] = 1
        ret['data']['msg'] = '参数错误！'
    except MsgError as e:
        ret['code'] = 1
        ret['data']['msg'] = e.msg
    except:
        ret['code'] = 1
        ret['data']['msg'] = '未知错误！'

    return jsonify(ret)

    
# 4. 修改指定payload

@bp.route('/payload/edit',methods=['POST'])
@verify
def edit_payload():
    data = request.get_json(force=True)
    ret = {'code': 0, 'data': {'msg': '修改成功！'}}
    try:
        payload = Payload.create_or_modify_by_json(data)
        tmp = payload.to_json()
        tmp.pop('value')
        update_base_info('payload','update',[tmp])
    except MsgError as e:
        ret['code'] = 1
        ret['data']['msg'] = e.msg
    except:
        ret['code'] = 1
        ret['data']['msg'] = '未知错误！'

    return jsonify(ret)

# 5. 获取所有项目基本信息

@bp.route('/project/list/all',methods=['GET'])
@verify
def get_all_project():
    ret = {'code': 0,'data':{'count': 0,'project_list':[]}}
    ret['data']['project_list'] = Project.get_all_project_info_with_json()
    ret['data']['count'] = len(ret['data']['project_list'])
    
    return jsonify(ret)

# 6. 获取指定项目详细内容

@bp.route('/project/<project_id>/get',methods=['GET'])
@verify
def get_project_detail(project_id):
    ret = {'code': 1, 'data': {'msg': '请求出错！'}}
    try:
        project_id = int(project_id)
        info = Project.get_project_detail_by_id(project_id)
        ret['code'] = 0
        ret['data'] = info
    except ValueError:
        ret['data']['msg'] = '参数错误！'
    except MsgError as e:
        ret['data']['msg'] = e.msg
    except:
        ret['code'] = 1
        ret['data']['msg'] = '未知错误！'
    return jsonify(ret)

# 7. 获取指定项目的攻击代码

@bp.route('/project/<project_id>/allcode',methods=['GET'])
@verify
def get_project_allcode(project_id):
    ret = {'code': 1, 'data': {}}
    try:
        project_id = int(project_id)
        info = Project.get_code_by_project_id(project_id)
        ret['code'] = 0
        ret['data']['project_id'] = project_id
        ret['data']['code'] = base64.b64encode(info.encode('utf8')).decode('utf8',errors='ignore')
    except ValueError:
        ret['data']['msg'] = '参数错误！'
    except MsgError as e:
        ret['data']['msg'] = e.msg
    except:
        ret['data']['msg'] = '未知错误！'
    return jsonify(ret)

# 8. 获取访问记录

@bp.route('/access_record',methods=['POST'])
@verify
def get_access_records():
    data = request.get_json(force=True)
    ret = {'code': 1, 'data': {}}
    try:
        records = Record.query_record(data)
    except:
        ret['data']['msg'] = '请求出错！'
        return ret
    ret['code'] = 0
    ret['data']['count'] = len(records)
    ret['data']['record_list'] = records
    return jsonify(ret)

# 9. 修改指定项目配置

@bp.route('/project/edit',methods=['POST'])
@verify
def edit_project():
    data = request.get_json(force=True)
    ret = {'code': 0, 'data': {'msg': '修改成功！'}}
    try:
        project = Project.create_or_modify_by_json(data)
        tmp = project.to_json()
        tmp['url'] = project.get_project_url()
        tmp['record_count'] = project.records.count()
        tmp.pop('config')
        update_base_info('project','update',[tmp])
    except MsgError as e:
        ret['code'] = 1
        ret['data']['msg'] = e.msg
    except:
        ret['code'] = 1
        ret['data']['msg'] = '未知错误！'

    return jsonify(ret)

# 10. 删除指定项目

@bp.route('/project/delete',methods=['POST'])
@verify
def delete_project():
    ret = {'code': 0, 'data': {'msg': '删除成功！'}}
    try:
        id_list = request.get_json(force=True)
        Project.delete_project_by_id_list(id_list)
        update_base_info('project','delete',[*id_list])
    except ValueError:
        ret['code'] = 1
        ret['data']['msg'] = '参数错误！'
    except MsgError as e:
        ret['code'] = 1
        ret['data']['msg'] = e.msg
    except:
        ret['code'] = 1
        ret['data']['msg'] = '未知错误！'
    return jsonify(ret)

# 11. 删除指定payload

@bp.route('/payload/delete',methods=['POST'])
@verify
def delete_payload():
    ret = {'code': 0, 'data': {'msg': '删除成功！'}}
    try:
        id_list = request.get_json(force=True)
        Payload.delete_payload_by_id_list(id_list)
        update_base_info('payload','delete',[*id_list])
    except ValueError:
        ret['code'] = 1
        ret['data']['msg'] = '参数错误！'
    except MsgError as e:
        ret['code'] = 1
        ret['data']['msg'] = e.msg
    except:
        ret['code'] = 1
        ret['data']['msg'] = '未知错误！'
    return jsonify(ret)


# 12. 创建一个新的payload

@bp.route('/payload/create',methods=['POST'])
@verify
def create_new_payload():
    data = request.get_json(force=True)
    ret = {'code': 0, 'data': {'msg': '创建成功！'}}
    try:
        payload = Payload.create_or_modify_by_json(data)
        tmp = payload.to_json()
        tmp.pop('value')
        update_base_info('payload','add',[tmp])
    except MsgError as e:
        ret['code'] = 1
        ret['data']['msg'] = e.msg
    except:
        ret['code'] = 1
        ret['data']['msg'] = '未知错误！'

    return jsonify(ret)


# 13. 创建一个新的项目

@bp.route('/project/create',methods=['POST'])
@verify
def create_project():
    ret = {'code': 0, 'data': {'msg': '创建成功！'}}
    data = request.get_json(force=True)
    try:
        project = Project.create_or_modify_by_json(data)
        tmp = project.to_json()
        tmp['url'] = project.get_project_url()
        tmp['record_count'] = project.records.count()
        tmp.pop('config')
        update_base_info('project','add',[tmp])
    except MsgError as e:
        ret['code'] = 1
        ret['data']['msg'] = e.msg
    except:
        ret['code'] = 1
        ret['data']['msg'] = '未知错误！'

    return jsonify(ret)


# 14. 新建一个分类

@bp.route('/category/create/<name>',methods=['GET'])
@verify
def create_category(name):
    ret = {'code': 0, 'data': {'msg': '创建成功！'}}
    try:
        name = urldecode(name)
        category = Category.create_category(name)
        update_base_info('category','add',[category.to_json()])
    except MsgError as e:
        ret['code'] = 1
        ret['data']['msg'] = e.msg
    except:
        ret['code'] = 1
        ret['data']['msg'] = '未知错误！'
    return jsonify(ret)

# 15. 列出所有分类
@bp.route('/category/all',methods=['GET'])
@verify 
def get_all_category():
    ret = {'code': 1, 'data': {}}
    info = Category.get_all_category_with_json()
    ret['code'] = 0
    ret['data']['count'] = len(info)
    ret['data']['category_list'] = info
    return jsonify(ret)

# 16. 删除指定分类
@bp.route('/category/<category_id>/delete',methods=['GET'])
@verify
def delete_category(category_id):
    ret = {'code': 1, 'data': {'msg': '请求出错！'}}
    try:
        category_id = int(category_id)
        Category.delete_category(category_id)
        ret['code'] = 0
        ret['data']['msg'] = '删除成功！'
        update_base_info('category','delete',[category_id])
    except ValueError:
        ret['data']['msg'] = '参数错误！'
    except MsgError as e:
        ret['data']['msg'] = e.msg
    except:
        ret['data']['msg'] = '未知错误！'
    return jsonify(ret)


# 17. 修改记录状态
@bp.route('/access_record/set_state',methods=['POST'])
@verify
def set_record_state():
    ret = {'code': 1, 'data': {'msg': '请求出错！'}}
    data = request.get_json(force=True)
    try:
        Record.set_state(data)
        ret['code'] = 0
        ret['data']['msg'] = '状态修改成功！'
    except MsgError as e:
        ret['data']['msg'] = e.msg
    except:
        ret['data']['msg'] = '未知错误！'
    return jsonify(ret)


# 18. 删除指定记录
@bp.route('/access_record/delete',methods=['POST'])
@verify
def delete_record():
    ret = {'code': 1, 'data': {'msg': '请求出错！'}}
    try:
        id_list = request.get_json(force=True)
        Record.delete_record_by_id_list(id_list)
        ret['code'] = 0
        ret['data']['msg'] = '删除成功！'
    except ValueError:
        ret['data']['msg'] = '参数错误！'
    except MsgError as e:
        ret['data']['msg'] = e.msg
    except:
        ret['data']['msg'] = '未知错误！'
    return jsonify(ret)

# 19. 重命名指定分类
@bp.route('/category/rename/<category_id>/<name>',methods=['GET'])
@verify
def rename_category(category_id,name):
    ret = {'code': 1, 'data': {'msg': '请求出错！'}}
    try:
        category_id = int(category_id)
        name = urldecode(name)
        category = Category.rename(category_id,name)
        ret['code'] = 0
        ret['data']['msg'] = '重命名成功！'
        update_base_info('category','update',[category.to_json()])
    except ValueError:
        ret['data']['msg'] = '参数错误！'
    except MsgError as e:
        ret['data']['msg'] = e.msg
    except:
        ret['data']['msg'] = '未知错误！'
    return jsonify(ret)

# 20. 获取基本配置信息
@bp.route('/config',methods=['GET'])
def get_base_config():
    ret = {'code': 0, 'data': {}}
    config = {
        'domain':current_app.config.get('WEB_ROOT'),
        'sessionName':current_app.config.get('SESSION_NAME')
    }
    ret["data"]["config"] = config
    return jsonify(ret)

# 21. 注销
@bp.route('/logout',methods=['GET'])
@verify
def logout():
    ret = {'code': 0, 'data': {'msg': '你已经注销！'}}
    cookie = request.cookies.get(current_app.config.get("SESSION_NAME"))
    if(not Session.session_destroy(cookie)):
        ret['code'] = 1
        ret['data']['msg'] = "请求出错！"
    return jsonify(ret)
